//
// Created by 37417 on 2024/4/5.
//

#ifndef HELLOTEST_OTHER_H
#define HELLOTEST_OTHER_H
namespace segment {
// Represents an aligned array of N elements of T. Data pointers can be
// reinterpreted as this type to generate vectorized loads/stores in a kernel.
template <typename T, int N>
class alignas(alignof(T) * N) AlignedVector {
 public:
  typedef T value_type;
  static constexpr const int kSize = N;

  AlignedVector() = default;

  // Uniform initialization.
  /*__host__ __device__*/ explicit AlignedVector(value_type uniform) {
    /*UNROLL_ON_DEVICE*/ for (int i = 0; i < kSize; ++i) {
      values_[i] = uniform;
    }
  }
  // Uniform initialization with explicit conversion.
  // Note: This is required for T=Eigen::half because it only supports explicit
  // conversions from other types and its template constructor is too relaxed
  // to be able to use std::is_constructible.
  template <typename U, typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
  /*__host__ __device__*/ explicit AlignedVector(U uniform_u) {
    value_type uniform(uniform_u);
    /*UNROLL_ON_DEVICE*/ for (int i = 0; i < kSize; ++i) {
      values_[i] = uniform;
    }
  }
  // Implicit conversion.
  template <typename U, typename std::enable_if<std::is_convertible<U, T>::value, int>::type = 0>
  /*__host__ __device__*/ AlignedVector(const AlignedVector<U, N>& other) {
    /*UNROLL_ON_DEVICE*/ for (int i = 0; i < kSize; ++i) {
      values_[i] = other[i];
    }
  }
  // Explicit conversion.
  template <typename U, typename std::enable_if<!std::is_convertible<U, T>::value &&
                                                    std::is_constructible<T, U>::value,
                                                int>::type = 0>
  /*__host__ __device__*/ explicit AlignedVector(const AlignedVector<U, N>& other) {
    /*UNROLL_ON_DEVICE*/ for (int i = 0; i < kSize; ++i) {
      values_[i] = T(other[i]);
    }
  }

  /*__host__ __device__*/ value_type& operator[](int i) { return values_[i]; }
  /*__host__ __device__*/ const value_type& operator[](int i) const { return values_[i]; }

#define DEFINE_BINARY_UPDATE_OPERATOR(op)                                         \
  /*__host__ __device__*/ AlignedVector& operator op(const AlignedVector & rhs) { \
    /*UNROLL_ON_DEVICE*/ for (int i = 0; i < kSize; ++i) {                        \
      values_[i] op rhs[i];                                                       \
    }                                                                             \
    return *this;                                                                 \
  }
  DEFINE_BINARY_UPDATE_OPERATOR(+=)
  DEFINE_BINARY_UPDATE_OPERATOR(-=)
  DEFINE_BINARY_UPDATE_OPERATOR(*=)
  DEFINE_BINARY_UPDATE_OPERATOR(/=)
#undef DEFINE_BINARY_UPDATE_OPERATOR

#define DEFINE_BINARY_OPERATOR(op)                                                     \
  friend /*__host__ __device__*/ AlignedVector operator op(const AlignedVector& lhs,   \
                                                           const AlignedVector& rhs) { \
    AlignedVector ret;                                                                 \
    /*UNROLL_ON_DEVICE*/ for (int i = 0; i < kSize; ++i) {                             \
      ret[i] = lhs[i] op rhs[i];                                                       \
    }                                                                                  \
    return ret;                                                                        \
  }
  DEFINE_BINARY_OPERATOR(+)
  DEFINE_BINARY_OPERATOR(-)
  DEFINE_BINARY_OPERATOR(*)
  DEFINE_BINARY_OPERATOR(/)
#undef DEFINE_BINARY_OPERATOR

#define DEFINE_BINARY_FUNCTION(func)                                            \
  friend /*__host__ __device__*/ AlignedVector func(const AlignedVector& lhs,   \
                                                    const AlignedVector& rhs) { \
    AlignedVector ret;                                                          \
    /*UNROLL_ON_DEVICE*/ for (int i = 0; i < kSize; ++i) {                      \
      ret[i] = func(lhs[i], rhs[i]);                                            \
    }                                                                           \
    return ret;                                                                 \
  }
  DEFINE_BINARY_FUNCTION(min)
  DEFINE_BINARY_FUNCTION(max)
#undef DEFINE_BINARY_FUNCTION

 private:
  value_type values_[N];
};
}  // namespace segment
#endif  // HELLOTEST_OTHER_H
